'''
https://leetcode.cn/problems/path-with-maximum-probability/description
'''
import heapq
from typing import List


class Solution:
    def maxProbability(self, n: int, edges: List[List[int]], succProb: List[float], start: int,
                       end: int) -> float:
        graph = [[] for _ in range(n)]
        for (u, v), w in zip(edges, succProb):
            graph[u].append((v, w))
            graph[v].append((u, w))

        distance = [.0] * n
        distance[start] = 1
        visited = [False] * n

        pq = [(-1, start)]
        while pq:
            prob, u = heapq.heappop(pq)
            if visited[u]: continue
            visited[u] = True
            for v, w in graph[u]:
                if not visited[v] and -prob * w > distance[v]:
                    distance[v] = -prob * w
                    heapq.heappush(pq, (-distance[v], v))
        # print(distance)
        return distance[end]

n = 3
edges = [[0,1],[1,2],[0,2]]
succProb = [0.5,0.5,0.2]
start = 0
end = 2
print(Solution().maxProbability(n, edges, succProb, start, end))


